import os
from langgraph.types import Command
from typing import TypedDict, Literal
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langgraph.constants import START, END
from langchain_core.prompts import ChatPromptTemplate
from src.agentic.agent.multi.StockAgent import StockAgent
from src.agentic.agent.multi.SearchAgent import SearchAgent
from src.agentic.agent.multi.WeatherAgent import WeatherAgent
from src.agentic.agent.multi.CalculateAgent import CalculateAgent

logger = getLogger()

class SupervisorState(TypedDict):
    query: str
    document: str
    answer: str

node_names = ["calculate", "weather", "stock", "search"]

class SupervisorAct(BaseModel):
    agent: Literal[*node_names, "finish"] = Field(description = "调用执行的下一步节点")

class SupervisorAgentic:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tool = { tool.name: tool for tool in agent_tools }

    def supervisor_node(self, state: SupervisorState) -> Command[Literal[*node_names, "__end__"]]:
        logger.info(f"SupervisorAgentic supervisor_node start")
        query = state["query"]
        answer = state.get("answer", None)
        logger.info(f"SupervisorAgentic supervisor_node answer: {answer}")
        route_template = """
            你是一位精准的问题-答案分析专家，请严格按照要求分析并解答，判断当前应归属的节点类型
            
            根据用户问题和答案，返回下列节点中的一个：
            calculate: 用于执行数学计算，只能用于数值计算
            weather: 用于查询城市当天及未来的实时天气情况，禁止用来查询过去时间的天气情况
            stock: 用于查询沪深300指数指定股票的开盘/收盘价
            search: 用于搜索真实世界的信息，如事实、人物、地点等，且不属于上述三类时使用
            finish: 答案不能为空，且满足以下任一条件时必须返回'finish'：
                1、答案包含与用户问题相关的关键词
                2、答案与用户问题语义含义一致
            请直接输出节点名称（如：search），不要附加任何解释或标点，禁止推测或扩展。

            用户问题： {question}
            答案：{answer}
        """
        route_prompt = ChatPromptTemplate.from_template(route_template)
        route_chain = route_prompt | self.llm_model.with_structured_output(SupervisorAct)
        route_result = route_chain.invoke({ "question": query, "answer": answer })
        logger.info(f"SupervisorAgentic supervisor_node route_result: {route_result}")
        goto = route_result.agent
        if goto == "finish":
            goto = END
        return Command(goto = route_result.agent, update = { "agent": goto })

    def calculate_node(self, state: SupervisorState) -> Command[Literal["supervisor"]]:
        logger.info(f"SupervisorAgentic calculate_node start")
        agent = CalculateAgent(self.llm_model, self.agent_tool["calculate_numerical"])
        response = agent.invoke(state["query"])
        return Command(goto = "supervisor", update = { "document": response["document"], "answer": response["answer"] })

    def weather_node(self, state: SupervisorState) -> Command[Literal["supervisor"]]:
        logger.info(f"SupervisorAgentic weather_node start")
        agent = WeatherAgent(self.llm_model, self.agent_tool["search_weather"])
        response = agent.invoke(state["query"])
        return Command(goto = "supervisor", update = { "document": response["document"], "answer": response["answer"] })

    def stock_node(self, state: SupervisorState) -> Command[Literal["supervisor"]]:
        logger.info(f"SupervisorAgentic stock_node start")
        agent = StockAgent(self.llm_model, self.agent_tool["snatch_stock_price"])
        response = agent.invoke(state["query"])
        return Command(goto = "supervisor", update = { "document": response["document"], "answer": response["answer"] })

    def search_node(self, state: SupervisorState) -> Command[Literal["supervisor"]]:
        logger.info(f"SupervisorAgentic search_node start")
        agent = SearchAgent(self.llm_model, self.agent_tool["search_web"])
        response = agent.invoke(state["query"])
        return Command(goto = "supervisor", update = { "document": response["document"], "answer": response["answer"] })

    def build_graph(self):
        logger.info(f"SupervisorAgentic build_graph start")
        graph = StateGraph(SupervisorState)
        graph.add_node("supervisor", self.supervisor_node)
        graph.add_node("calculate", self.calculate_node)
        graph.add_node("weather", self.weather_node)
        graph.add_node("stock", self.stock_node)
        graph.add_node("search", self.search_node)

        graph.add_edge(START, "supervisor")

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agentic_Supervisor_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"SupervisorAgentic invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        answer = response.get("answer", None)
        logger.info(f"SupervisorAgentic invoke answer len: {len(answer)}")
        return { "retrieve_docs": response["document"], "chain_result": answer }
